import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt


# 1. 生成正弦波数据（仅使用 PyTorch）
def generate_sine_wave(seq_length=10, num_samples=1000):
    x = torch.linspace(0, 100, num_samples)  # 生成 1000 个等间距数据点
    y = torch.sin(x)  # 计算正弦值

    X_data, Y_data = [], []
    for i in range(len(y) - seq_length):
        X_data.append(y[i:i + seq_length].unsqueeze(-1))  # 过去 seq_length 作为输入
        Y_data.append(y[i + seq_length])  # 预测下一个点

    return torch.stack(X_data), torch.tensor(Y_data).unsqueeze(-1)


# 生成数据
seq_length = 10  # 序列长度
X, Y = generate_sine_wave(seq_length)

# 划分训练集和测试集
train_size = int(0.8 * len(X))
X_train, X_test = X[:train_size], X[train_size:]
Y_train, Y_test = Y[:train_size], Y[train_size:]


# 2. 定义 RNN 模型
class SimpleRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=1):
        super(SimpleRNN, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers

        self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size)  # 初始化隐藏状态
        out, _ = self.rnn(x, h0)
        out = self.fc(out[:, -1, :])  # 取最后一个时间步的输出
        return out


# 3. 训练模型
# 超参数
input_size = 1
hidden_size = 32
output_size = 1
num_layers = 1
num_epochs = 100
learning_rate = 0.001

# 初始化模型
model = SimpleRNN(input_size, hidden_size, output_size, num_layers)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 训练
for epoch in range(num_epochs):
    model.train()
    optimizer.zero_grad()

    outputs = model(X_train)
    loss = criterion(outputs, Y_train)

    loss.backward()
    optimizer.step()

    if (epoch + 1) % 10 == 0:
        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 4. 评估与绘图
model.eval()
with torch.no_grad():
    predictions = model(X_test)

# 画图
plt.figure(figsize=(10, 5))
plt.plot(Y_test.numpy(), label="Real Data")
plt.plot(predictions.numpy(), label="Predicted Data")
plt.legend()
plt.title("RNN Sine Wave Prediction")
plt.show()
